import argparse
import json
import os
import os.path as osp
import re
from collections import OrderedDict
from datetime import datetime
from pathlib import Path
from zipfile import ZipFile

import yaml
from modelindex.load_model_index import load
from rich.console import Console
from rich.syntax import Syntax
from rich.table import Table

console = Console()
MMSELFSUP_ROOT = Path(__file__).absolute().parents[2]
CYCLE_LEVELS = ['month', 'quarter', 'half-year', 'no-training']
METRICS_MAP = {
    'Top 1 Accuracy': 'accuracy/top1',
    'Top 5 Accuracy': 'accuracy/top5'
}


class RangeAction(argparse.Action):

    def __call__(self, parser, namespace, values: str, option_string):
        matches = re.match(r'([><=]*)([-\w]+)', values)
        if matches is None:
            raise ValueError(f'Unavailable range option {values}')
        symbol, range_str = matches.groups()
        assert range_str in CYCLE_LEVELS, \
            f'{range_str} are not in {CYCLE_LEVELS}.'
        level = CYCLE_LEVELS.index(range_str)
        symbol = symbol or '<='
        ranges = set()
        if '=' in symbol:
            ranges.add(level)
        if '>' in symbol:
            ranges.update(range(level + 1, len(CYCLE_LEVELS)))
        if '<' in symbol:
            ranges.update(range(level))
        assert len(ranges) > 0, 'No range are selected.'
        setattr(namespace, self.dest, ranges)


def parse_args():
    parser = argparse.ArgumentParser(
        description='Train models (in models.yml) and compare accuracy.')
    parser.add_argument(
        'partition', type=str, help='Cluster partition to use.')
    parser.add_argument(
        '--job-name',
        type=str,
        default='selfsup-benchmark',
        help='Slurm job name prefix')
    parser.add_argument('--port', type=int, default=29777, help='dist port')
    parser.add_argument(
        '--models', nargs='+', type=str, help='Specify model names to run.')
    parser.add_argument(
        '--range',
        type=str,
        default={0},
        action=RangeAction,
        metavar='{month,quarter,half-year,no-training}',
        help='The training benchmark range, "no-training" means all models '
        "including those we haven't trained.")
    parser.add_argument(
        '--work-dir',
        default='work_dirs/benchmark_pretrain_cls',
        help='the dir to save train log')
    parser.add_argument(
        '--run', action='store_true', help='run script directly')
    parser.add_argument(
        '--local',
        action='store_true',
        help='run at local instead of cluster.')
    parser.add_argument(
        '--mail', type=str, help='Mail address to watch train status.')
    parser.add_argument(
        '--mail-type',
        nargs='+',
        default=['BEGIN', 'END', 'FAIL'],
        choices=['NONE', 'BEGIN', 'END', 'FAIL', 'REQUEUE', 'ALL'],
        help='Mail address to watch train status.')
    parser.add_argument(
        '--quotatype',
        default=None,
        choices=['reserved', 'auto', 'spot'],
        help='Quota type, only available for phoenix-slurm>=0.2')
    parser.add_argument(
        '--summary',
        action='store_true',
        help='Summarize benchmark train results.')
    parser.add_argument(
        '--save',
        action='store_true',
        help='Save the summary and archive log files.')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        type=str,
        default=[],
        help='Config options for all config files.')

    args = parser.parse_args()
    return args


def get_gpu_number(model_info):
    config = osp.basename(model_info.config)
    matches = re.match(r'.*[-_](\d+)xb(\d+).*', config)
    if matches is None:
        raise ValueError(
            f'Cannot get gpu numbers from the config name {config}')
    gpus = int(matches.groups()[0])
    return gpus


def get_pretrain_epoch(model_info):
    config = osp.basename(model_info.config)
    matches = re.match(r'.*[-_](\d+)e[-_].*', config)
    if matches is None:
        raise ValueError(
            f'Cannot get epoch setting from the config name {config}')
    epoch = int(matches.groups()[0])
    return epoch


def create_train_job_batch(commands, model_info, args, port, script_name):

    fname = model_info.name

    gpus = get_gpu_number(model_info)
    gpus_per_node = min(gpus, 8)

    config = Path(model_info.config)
    assert config.exists(), f'"{fname}": {config} not found.'

    work_dir = Path(args.work_dir) / fname
    work_dir.mkdir(parents=True, exist_ok=True)

    if args.quotatype is not None:
        quota_cfg = f'--quotatype {args.quotatype} '
    else:
        quota_cfg = ''

    launcher = 'none' if args.local else 'slurm'
    job_name = f'{args.job_name}_{fname}'
    job_script = (f'#!/bin/bash\n'
                  f'srun -p {args.partition} '
                  f'--job-name {job_name} '
                  f'--gres=gpu:{gpus_per_node} '
                  f'{quota_cfg}'
                  f'--ntasks-per-node={gpus_per_node} '
                  f'--ntasks={gpus} '
                  f'--cpus-per-task=12 '
                  f'--kill-on-bad-exit=1 '
                  f'python -u {script_name} {config} '
                  f'--work-dir={work_dir} '
                  f'--cfg-option env_cfg.dist_cfg.port={port} '
                  f'{" ".join(args.cfg_options)} '
                  f'default_hooks.checkpoint.max_keep_ckpts=1 '
                  f'--launcher={launcher}\n')

    commands.append(f'echo "{config}"')

    # downstream classification task
    cls_config = None
    task = getattr(model_info, 'task', None)
    if task is not None:
        for downstream in model_info.data['Downstream']:
            if task == downstream['Results'][0]['Task']:
                cls_config = downstream['Config']
                break
    else:
        cls_config = None

    if cls_config:
        fname = model_info.name

        gpus = get_gpu_number(model_info)
        gpus_per_node = min(gpus, 8)

        cls_config_path = Path(cls_config)
        assert cls_config_path.exists(), f'"{fname}": {cls_config} not found.'

        job_name = f'{args.job_name}_{fname}'

        cls_work_dir = work_dir / Path(
            cls_config.split('/')[-1].replace('.py', ''))
        cls_work_dir.mkdir(parents=True, exist_ok=True)

        srun_args = ''
        if args.quotatype is not None:
            srun_args = srun_args.join(f'--quotatype {args.quotatype}')

        # get pretrain weights path
        epoch = get_pretrain_epoch(model_info)
        ckpt = work_dir / f'epoch_{epoch}.pth'

        launcher = 'none' if args.local else 'slurm'
        cls_job_script = (
            f'\n'
            f'mim train mmcls {cls_config} '
            f'--launcher {launcher} '
            f'-G {gpus} '
            f'--gpus-per-node {gpus_per_node} '
            f'--cpus-per-task 12 '
            f'--partition {args.partition} '
            f'--srun-args "{srun_args}" '
            f'--work-dir {cls_work_dir} '
            f'--cfg-option model.backbone.init_cfg.type=Pretrained '
            f'model.backbone.init_cfg.checkpoint={ckpt} '
            f'model.backbone.init_cfg.prefix=backbone. '
            f'default_hooks.checkpoint.max_keep_ckpts=1 '
            f'{" ".join(args.cfg_options)}\n')

        commands.append(f'echo "{cls_config}"')

    with open(work_dir / 'job.sh', 'w') as f:
        f.write(job_script)
        if cls_config:
            f.write(cls_job_script)

    commands.append(
        f'nohup bash {work_dir}/job.sh > {work_dir}/out.log 2>&1 &')

    return work_dir / 'job.sh'


def train(models, args):
    script_name = osp.join('tools', 'train.py')
    port = args.port

    commands = []

    for model_info in models.values():
        script_path = create_train_job_batch(commands, model_info, args, port,
                                             script_name)
        port += 1

    command_str = '\n'.join(commands)

    preview = Table()
    preview.add_column(str(script_path))
    preview.add_column('Shell command preview')
    preview.add_row(
        Syntax.from_path(
            script_path,
            background_color='default',
            line_numbers=True,
            word_wrap=True),
        Syntax(
            command_str,
            'bash',
            background_color='default',
            line_numbers=True,
            word_wrap=True))
    console.print(preview)

    if args.run:
        os.system(command_str)
    else:
        console.print('Please set "--run" to start the job')


def save_summary(summary_data, models_map, work_dir):
    date = datetime.now().strftime('%Y%m%d-%H%M%S')
    zip_path = work_dir / f'archive-{date}.zip'
    zip_file = ZipFile(zip_path, 'w')
    summary_path = work_dir / 'benchmark_summary.md'
    file = open(summary_path, 'w')
    headers = [
        'Model', 'Top-1 Expected(%)', 'Top-1 (%)', 'Top-1 best(%)',
        'best epoch', 'Top-5 Expected (%)', 'Top-5 (%)', 'Config', 'Log'
    ]
    file.write('# Train Benchmark Regression Summary\n')
    file.write('| ' + ' | '.join(headers) + ' |\n')
    file.write('|:' + ':|:'.join(['---'] * len(headers)) + ':|\n')
    for model_name, summary in summary_data.items():
        if len(summary) == 0:
            # Skip models without results
            continue
        row = [model_name]
        if 'Top 1 Accuracy' in summary:
            metric = summary['Top 1 Accuracy']
            row.append(f"{metric['expect']:.2f}")
            row.append(f"{metric['last']:.2f}")
            row.append(f"{metric['best']:.2f}")
            row.append(f"{metric['best_epoch']:.2f}")
        else:
            row.extend([''] * 4)
        if 'Top 5 Accuracy' in summary:
            metric = summary['Top 5 Accuracy']
            row.append(f"{metric['expect']:.2f}")
            row.append(f"{metric['last']:.2f}")
        else:
            row.extend([''] * 2)

        model_info = models_map[model_name]
        row.append(model_info.config)
        row.append(str(summary['log_file'].relative_to(work_dir)))
        zip_file.write(summary['log_file'])
        file.write('| ' + ' | '.join(row) + ' |\n')
    file.close()
    zip_file.write(summary_path)
    zip_file.close()
    print('Summary file saved at ' + str(summary_path))
    print('Log files archived at ' + str(zip_path))


def show_summary(summary_data):
    table = Table(title='Train Benchmark Regression Summary')
    table.add_column('Model')
    for metric in METRICS_MAP:
        table.add_column(f'{metric} (expect)')
        table.add_column(f'{metric}')
        table.add_column(f'{metric} (best)')

    def set_color(value, expect):
        if value > expect:
            return 'green'
        elif value > expect - 0.2:
            return 'white'
        else:
            return 'red'

    for model_name, summary in summary_data.items():
        row = [model_name]
        for metric_key in METRICS_MAP:
            if metric_key in summary:
                metric = summary[metric_key]
                expect = metric['expect']
                last = metric['last']
                last_epoch = metric['last_epoch']
                last_color = set_color(last, expect)
                best = metric['best']
                best_color = set_color(best, expect)
                best_epoch = metric['best_epoch']
                row.append(f'{expect:.2f}')
                row.append(
                    f'[{last_color}]{last:.2f}[/{last_color}] ({last_epoch})')
                row.append(
                    f'[{best_color}]{best:.2f}[/{best_color}] ({best_epoch})')
        table.add_row(*row)

    console.print(table)


def summary(models, args):

    work_dir = Path(args.work_dir)
    dir_map = {p.name: p for p in work_dir.iterdir() if p.is_dir()}

    summary_data = {}
    for model_name, model_info in models.items():

        summary_data[model_name] = {}

        if model_name not in dir_map:
            continue

        # Skip if not found any vis_data folder.
        sub_dir = dir_map[model_name]
        log_files = [f for f in sub_dir.glob('*/*/vis_data/scalars.json')]
        if len(log_files) == 0:
            continue
        log_file = sorted(log_files)[-1]

        # parse train log
        with open(log_file) as f:
            json_logs = [json.loads(s) for s in f.readlines()]
            val_logs = [
                log for log in json_logs
                # TODO: need a better method to extract validate log
                if 'loss' not in log and 'accuracy/top1' in log
            ]

        if len(val_logs) == 0:
            continue

        for downstream in model_info.data['Downstream']:
            if model_info.task == downstream['Results'][0]['Task']:
                expect_metrics = downstream['Results'][0]['Metrics']
                break

        # extract metrics
        summary = {'log_file': log_file}
        for key_yml, key_res in METRICS_MAP.items():
            if key_yml in expect_metrics:
                assert key_res in val_logs[-1], \
                    f'{model_name}: No metric "{key_res}"'
                expect_result = float(expect_metrics[key_yml])
                last = float(val_logs[-1][key_res])
                best_log, best_epoch = sorted(
                    zip(val_logs, range(len(val_logs))),
                    key=lambda x: x[0][key_res])[-1]
                best = float(best_log[key_res])

                summary[key_yml] = dict(
                    expect=expect_result,
                    last=last,
                    last_epoch=len(val_logs),
                    best=best,
                    best_epoch=best_epoch + 1)
        summary_data[model_name].update(summary)

    show_summary(summary_data)
    if args.save:
        save_summary(summary_data, models, work_dir)


def main():
    args = parse_args()

    model_index_file = MMSELFSUP_ROOT / 'model-index.yml'
    model_index = load(str(model_index_file))
    model_index.build_models_with_collections()
    all_models = {model.name: model for model in model_index.models}

    with open(Path(__file__).parent / 'models.yml', 'r') as f:
        train_items = yaml.safe_load(f)
    models = OrderedDict()
    for item in train_items:
        name = item['Name']
        model_info = all_models[item['Name']]
        model_info.cycle = item.get('Cycle', None)
        model_info.task = item.get('Task', None)
        cycle = getattr(model_info, 'cycle', 'month')
        cycle_level = CYCLE_LEVELS.index(cycle)
        if cycle_level in args.range:
            models[name] = model_info

    if args.models:
        patterns = [re.compile(pattern) for pattern in args.models]
        filter_models = {}
        for k, v in models.items():
            if any([re.match(pattern, k) for pattern in patterns]):
                filter_models[k] = v
        if len(filter_models) == 0:
            print('No model found, please specify models in:')
            print('\n'.join(models.keys()))
            return
        models = filter_models

    if args.summary:
        summary(models, args)
    else:
        train(models, args)


if __name__ == '__main__':
    main()
